package com.eqsoft.gesturerecognition;

import java.util.Vector;

public class HiddenMarkovModel {
	private int numStates, numObservations;
	private double pi[];
	private double A[][];
	private double B[][];
	
	public HiddenMarkovModel(int numStates, int numObservations)
	{
		this.numStates = numStates;
		this.numObservations = numObservations;
		
		pi = new double[numStates];
		A = new double[numStates][numStates];
		B = new double[numStates][numObservations];
		init();
	}
	
	private void init()
	{
		pi[0] = 1;
		for (int i = 1; i < numStates; i++)
			pi[i] = 0;
		
		int delta = 2;
		for (int i = 0; i < numStates; i++)
			for (int j = 0; j < numStates; j++)
			{
				if (j < i)
					A[i][j] = 0;
				else if (i <= j && i + delta >= j)
					A[i][j] = Math.max(1d/(delta+1), 1d/(numStates - i));
			}
		
		for (int i = 0; i < numStates; i++)
			for (int j = 0; j < numObservations; j++)
				B[i][j] = 1d/numObservations;
	}
	
	private double[][] forward(int[] o)
	{
		double[][] fwd = new double[numStates][o.length];
		for (int i = 0; i < fwd.length; i++)
			fwd[i][0] = pi[i] * B[i][o[0]];
		
		for (int i = 1; i < o.length; i++)
			for (int j = 0; j < fwd.length; j++)
			{
				double sum = 0;
				for (int k = 0; k < numStates; k++)
					sum += fwd[k][i-1] * A[k][j];
				fwd[j][i] = sum * B[j][o[i]];
			}
		
		
		return fwd;
	}
	private double[][] backward(int[] o)
	{
		double[][] bwd = new double[numStates][o.length];
		for (int i = 0; i < numStates; i++)
			bwd[i][o.length - 1] = 1;
		
		for (int i = o.length - 2; i >= 0; i--)
			for (int j = 0; j < numStates; j++)
			{
				bwd[j][i] = 0;
				for (int k = 0; k < numStates; k++)
					bwd[j][i] += bwd[k][i+1] * A[j][k] * B[k][o[i+1]];
			}
		
		return bwd;
	}
	
	public double getProb(int[] o)
	{
		double p = 0;
		double[][] fwd = forward(o);
		for (int i = 0; i < numStates; i++)
			p += fwd[i][numStates - 1];
		return p;
	}
	
	public void train(Vector<int[]> Q)
	{
		double[][] a = new double[numStates][numStates];
		double[][] b = new double[numStates][numObservations];
		
		for (int i = 0; i < numStates; i++)
			for (int j = 0; j < numStates; j++)
			{
				double num = 0, den = 0;
				
				for (int k = 0; k < Q.size(); k++)
				{
					int[] seq = Q.get(k);
					
					double[][] fwd = forward(seq);
					double[][] bwd = backward(seq);
					double prob = getProb(seq);
					double num_i = 0, den_i = 0;
					
					for (int l = 0; l < seq.length - 1; l++)
					{
						num_i += fwd[i][l] * A[i][j] * B[j][seq[l+1]] *bwd[j][l+1];
						den_i += fwd[i][l] * bwd[i][l];
					}
					
					num += (1/prob) * num_i;
					den += (1/prob) * den_i;
				}
				
				a[i][j] = num / den;
			}
		
		for (int i = 0; i < numStates; i++)
			for (int j = 0; j < numObservations; j++)
			{
				double num = 0, den = 0;
				for (int k = 0; k < Q.size(); k++)
				{
					int[] seq = Q.get(k);
					
					double[][] fwd = forward(seq);
					double[][] bwd = backward(seq);
					double prob = getProb(seq);
					double num_i = 0, den_i = 0;
					
					for (int l = 0; l < seq.length - 1; l++)
					{
						if (seq[l] == j)
							num_i += fwd[i][l] * bwd[i][l];
						den_i += fwd[i][l] * bwd[i][l];
					}
					
					num += (1/prob) * num_i;
					den += (1/prob) * den_i;
				}
				
				b[i][j] = num / den;
			}
		
		A = a;
		B = b;
	}
}
